from glob import glob
from multiprocessing import Pool

import pandas as pd
from tqdm import tqdm

def to_pickle(one_path,vv):
    for path in one_path:

        one = pd.read_pickle(path, compression="zip")
        new_data_list = []
        one=list(one.values())
        while len(one):
            v=one.pop()
            # new_voc=voc[voc["voc"].isin(v)]
            v = ["<|start|>"] + v + ["<|end|>"]
            v = pd.DataFrame({0:v[:-2], 1:v[1:-1], 2:v[2:]})
            v = v[~v[1].isin(vv)]
            v=v.drop_duplicates([0, 1, 2]).values.tolist()
            new_data_list += v
        pd.to_pickle(new_data_list, path.replace("skypile_jieba", "sky_pile_thr"), compression="zip")

if __name__ == '__main__':

    paths = glob("F:/skypile_jieba/*")
    pool = Pool(processes=20)
    p_list=[]

    voc = pd.read_pickle("total_voc_new.pkl")
    voc = voc["voc"]

    for i in tqdm(range(0,len(paths),180)):


        # to_pickle(paths[i:i+80],)
        p = pool.apply_async(func=to_pickle, args=(paths[i:i+180],voc))
        p_list.append(p)
    for p in tqdm(p_list):
        p.get()
    pool.close()
    pool.join()
